package redis

import (
    "context"
    "encoding/json"
    "fmt"
    "strings"
    "time"

    "github.com/go-redis/redis/v8"
    "github.com/yangyongzhen/mqtt-broker/internal/protocol/common"
    "github.com/yangyongzhen/mqtt-broker/internal/storage"
)

type RedisStore struct {
    client *redis.Client
    ctx    context.Context
}

type Config struct {
    Address        string        `json:"address" yaml:"address"`
    Password       string        `json:"password" yaml:"password"`
    DB             int           `json:"db" yaml:"db"`
    PoolSize       int           `json:"pool_size" yaml:"pool_size"`
    MinIdleConn    int           `json:"min_idle_conn" yaml:"min_idle_conn"`
    MaxRetries     int           `json:"max_retries" yaml:"max_retries"`
    DialTimeout    time.Duration `json:"dial_timeout" yaml:"dial_timeout"`
    ReadTimeout    time.Duration `json:"read_timeout" yaml:"read_timeout"`
    WriteTimeout   time.Duration `json:"write_timeout" yaml:"write_timeout"`
    SessionTTL     time.Duration `json:"session_ttl" yaml:"session_ttl"`
    MessageTTL     time.Duration `json:"message_ttl" yaml:"message_ttl"`
    MaxMessages    int64         `json:"max_messages" yaml:"max_messages"`
}

func NewRedisStore(config *Config) (*RedisStore, error) {
    // 设置默认值
    if config.PoolSize == 0 {
        config.PoolSize = 100
    }
    if config.MinIdleConn == 0 {
        config.MinIdleConn = 10
    }
    if config.MaxRetries == 0 {
        config.MaxRetries = 3
    }
    if config.DialTimeout == 0 {
        config.DialTimeout = 10 * time.Second
    }
    if config.ReadTimeout == 0 {
        config.ReadTimeout = 3 * time.Second
    }
    if config.WriteTimeout == 0 {
        config.WriteTimeout = 3 * time.Second
    }
    if config.SessionTTL == 0 {
        config.SessionTTL = 24 * time.Hour
    }
    if config.MessageTTL == 0 {
        config.MessageTTL = 24 * time.Hour
    }
    if config.MaxMessages == 0 {
        config.MaxMessages = 1000
    }

    rdb := redis.NewClient(&redis.Options{
        Addr:         config.Address,
        Password:     config.Password,
        DB:           config.DB,
        PoolSize:     config.PoolSize,
        MinIdleConns: config.MinIdleConn,
        MaxRetries:   config.MaxRetries,
        DialTimeout:  config.DialTimeout,
        ReadTimeout:  config.ReadTimeout,
        WriteTimeout: config.WriteTimeout,
    })

    ctx := context.Background()
    
    // 测试连接
    if err := rdb.Ping(ctx).Err(); err != nil {
        return nil, fmt.Errorf("failed to connect to Redis: %v", err)
    }

    return &RedisStore{
        client: rdb,
        ctx:    ctx,
    }, nil
}

// Redis键名常量
const (
    sessionPrefix      = "mqtt:session:"
    messagePrefix      = "mqtt:messages:"
    retainedPrefix     = "mqtt:retained:"
    subscriptionPrefix = "mqtt:subscriptions:"
    retainedIndex      = "mqtt:retained_index"
)

// 会话管理
func (r *RedisStore) SaveSession(clientID string, session *storage.Session) error {
    key := sessionPrefix + clientID
    
    // 序列化会话数据
    sessionData := map[string]interface{}{
        "client_id":     session.ClientID,
        "clean_session": session.CleanSession,
        "last_seen":     session.LastSeen.Unix(),
    }
    
    // 使用事务保存会话和订阅
    pipe := r.client.TxPipeline()
    
    // 保存基本会话信息
    pipe.HMSet(r.ctx, key, sessionData)
    pipe.Expire(r.ctx, key, 24*time.Hour)
    
    // 保存订阅信息
    subsKey := subscriptionPrefix + clientID
    pipe.Del(r.ctx, subsKey) // 先清空现有订阅
    
    if len(session.Subscriptions) > 0 {
        subsData := make(map[string]interface{})
        for topic, sub := range session.Subscriptions {
            subBytes, _ := json.Marshal(sub)
            subsData[topic] = string(subBytes)
        }
        pipe.HMSet(r.ctx, subsKey, subsData)
        pipe.Expire(r.ctx, subsKey, 24*time.Hour)
    }
    
    // 保存待确认消息
    if len(session.PendingAcks) > 0 {
        acksKey := key + ":pending_acks"
        pipe.Del(r.ctx, acksKey)
        
        acksData := make(map[string]interface{})
        for packetID, msg := range session.PendingAcks {
            msgBytes, _ := json.Marshal(msg)
            acksData[fmt.Sprintf("%d", packetID)] = string(msgBytes)
        }
        pipe.HMSet(r.ctx, acksKey, acksData)
        pipe.Expire(r.ctx, acksKey, 24*time.Hour)
    }
    
    _, err := pipe.Exec(r.ctx)
    return err
}

func (r *RedisStore) LoadSession(clientID string) (*storage.Session, error) {
    key := sessionPrefix + clientID
    
    // 检查会话是否存在
    exists, err := r.client.Exists(r.ctx, key).Result()
    if err != nil {
        return nil, fmt.Errorf("failed to check session existence: %v", err)
    }
    if exists == 0 {
        return nil, nil
    }
    
    // 获取基本会话信息
    sessionData, err := r.client.HGetAll(r.ctx, key).Result()
    if err != nil {
        return nil, fmt.Errorf("failed to load session: %v", err)
    }
    
    session := &storage.Session{
        ClientID:      clientID,
        Subscriptions: make(map[string]*common.Subscription),
        PendingAcks:   make(map[uint16]*common.Message),
    }
    
    if cleanSession, ok := sessionData["clean_session"]; ok {
        session.CleanSession = cleanSession == "true"
    }
    
    if lastSeen, ok := sessionData["last_seen"]; ok {
        if timestamp, err := time.Parse("2006-01-02 15:04:05", lastSeen); err == nil {
            session.LastSeen = timestamp
        }
    }
    
    // 加载订阅信息
    subsKey := subscriptionPrefix + clientID
    subsData, err := r.client.HGetAll(r.ctx, subsKey).Result()
    if err == nil {
        for topic, subStr := range subsData {
            var sub common.Subscription
            if json.Unmarshal([]byte(subStr), &sub) == nil {
                session.Subscriptions[topic] = &sub
            }
        }
    }
    
    // 加载待确认消息
    acksKey := key + ":pending_acks"
    acksData, err := r.client.HGetAll(r.ctx, acksKey).Result()
    if err == nil {
        for packetIDStr, msgStr := range acksData {
            var msg common.Message
            if json.Unmarshal([]byte(msgStr), &msg) == nil {
                if packetID := parsePacketID(packetIDStr); packetID > 0 {
                    session.PendingAcks[packetID] = &msg
                }
            }
        }
    }
    
    return session, nil
}

func (r *RedisStore) DeleteSession(clientID string) error {
    keys := []string{
        sessionPrefix + clientID,
        messagePrefix + clientID,
        subscriptionPrefix + clientID,
        sessionPrefix + clientID + ":pending_acks",
    }
    
    return r.client.Del(r.ctx, keys...).Err()
}

// 消息存储
func (r *RedisStore) SaveMessage(clientID string, message *common.Message) error {
    key := messagePrefix + clientID
    
    msgBytes, err := json.Marshal(message)
    if err != nil {
        return fmt.Errorf("failed to marshal message: %v", err)
    }
    
    // 使用有序集合存储消息，以时间戳为分数
    score := float64(message.Timestamp.UnixNano())
    member := string(msgBytes)
    
    pipe := r.client.Pipeline()
    pipe.ZAdd(r.ctx, key, &redis.Z{Score: score, Member: member})
    pipe.ZRemRangeByRank(r.ctx, key, 0, -1001) // 保持最新1000条消息
    pipe.Expire(r.ctx, key, 24*time.Hour)
    
    _, err = pipe.Exec(r.ctx)
    return err
}

func (r *RedisStore) LoadMessages(clientID string) ([]*common.Message, error) {
    key := messagePrefix + clientID
    
    // 按时间戳倒序获取消息
    data, err := r.client.ZRevRange(r.ctx, key, 0, -1).Result()
    if err != nil {
        return nil, fmt.Errorf("failed to load messages: %v", err)
    }
    
    messages := make([]*common.Message, 0, len(data))
    for _, item := range data {
        var message common.Message
        if err := json.Unmarshal([]byte(item), &message); err != nil {
            continue // 跳过无效消息
        }
        messages = append(messages, &message)
    }
    
    return messages, nil
}

func (r *RedisStore) DeleteMessage(clientID string, packetID uint16) error {
    key := messagePrefix + clientID
    
    // 获取所有消息，找到匹配的删除
    data, err := r.client.ZRange(r.ctx, key, 0, -1).Result()
    if err != nil {
        return err
    }
    
    for _, item := range data {
        var message common.Message
        if err := json.Unmarshal([]byte(item), &message); err != nil {
            continue
        }
        
        if message.PacketID == packetID {
            return r.client.ZRem(r.ctx, key, item).Err()
        }
    }
    
    return nil
}

// 保留消息
func (r *RedisStore) SaveRetainedMessage(topic string, message *common.Message) error {
    key := retainedPrefix + topic
    
    if message == nil || len(message.Payload) == 0 {
        // 删除保留消息
        pipe := r.client.Pipeline()
        pipe.Del(r.ctx, key)
        pipe.SRem(r.ctx, retainedIndex, topic)
        _, err := pipe.Exec(r.ctx)
        return err
    }
    
    msgBytes, err := json.Marshal(message)
    if err != nil {
        return fmt.Errorf("failed to marshal retained message: %v", err)
    }
    
    // 保存消息并添加到索引
    pipe := r.client.Pipeline()
    pipe.Set(r.ctx, key, msgBytes, 0) // 保留消息不过期
    pipe.SAdd(r.ctx, retainedIndex, topic)
    
    _, err = pipe.Exec(r.ctx)
    return err
}

func (r *RedisStore) LoadRetainedMessage(topic string) (*common.Message, error) {
    key := retainedPrefix + topic
    
    data, err := r.client.Get(r.ctx, key).Result()
    if err != nil {
        if err == redis.Nil {
            return nil, nil
        }
        return nil, fmt.Errorf("failed to load retained message: %v", err)
    }
    
    var message common.Message
    if err := json.Unmarshal([]byte(data), &message); err != nil {
        return nil, fmt.Errorf("failed to unmarshal retained message: %v", err)
    }
    
    return &message, nil
}

func (r *RedisStore) DeleteRetainedMessage(topic string) error {
    key := retainedPrefix + topic
    
    pipe := r.client.Pipeline()
    pipe.Del(r.ctx, key)
    pipe.SRem(r.ctx, retainedIndex, topic)
    
    _, err := pipe.Exec(r.ctx)
    return err
}

// 订阅管理
func (r *RedisStore) SaveSubscription(clientID string, subscription *common.Subscription) error {
    key := subscriptionPrefix + clientID
    
    subBytes, err := json.Marshal(subscription)
    if err != nil {
        return fmt.Errorf("failed to marshal subscription: %v", err)
    }
    
    pipe := r.client.Pipeline()
    pipe.HSet(r.ctx, key, subscription.Topic, subBytes)
    pipe.Expire(r.ctx, key, 24*time.Hour)
    
    _, err = pipe.Exec(r.ctx)
    return err
}

func (r *RedisStore) LoadSubscriptions(clientID string) ([]*common.Subscription, error) {
    key := subscriptionPrefix + clientID
    
    data, err := r.client.HGetAll(r.ctx, key).Result()
    if err != nil {
        return nil, fmt.Errorf("failed to load subscriptions: %v", err)
    }
    
    subscriptions := make([]*common.Subscription, 0, len(data))
    for _, item := range data {
        var subscription common.Subscription
        if err := json.Unmarshal([]byte(item), &subscription); err != nil {
            continue
        }
        subscriptions = append(subscriptions, &subscription)
    }
    
    return subscriptions, nil
}

func (r *RedisStore) DeleteSubscription(clientID string, topic string) error {
    key := subscriptionPrefix + clientID
    return r.client.HDel(r.ctx, key, topic).Err()
}

// 扩展方法：获取所有保留消息（类似内存存储）
func (r *RedisStore) GetAllRetainedMessages() (map[string]*common.Message, error) {
    topics, err := r.client.SMembers(r.ctx, retainedIndex).Result()
    if err != nil {
        return nil, err
    }
    
    result := make(map[string]*common.Message)
    
    if len(topics) == 0 {
        return result, nil
    }
    
    // 批量获取保留消息
    keys := make([]string, len(topics))
    for i, topic := range topics {
        keys[i] = retainedPrefix + topic
    }
    
    values, err := r.client.MGet(r.ctx, keys...).Result()
    if err != nil {
        return nil, err
    }
    
    for i, value := range values {
        if value == nil {
            continue
        }
        
        var message common.Message
        if err := json.Unmarshal([]byte(value.(string)), &message); err != nil {
            continue
        }
        
        result[topics[i]] = &message
    }
    
    return result, nil
}

// 扩展方法：获取匹配主题过滤器的保留消息
func (r *RedisStore) GetMatchingRetainedMessages(topicFilter string) ([]*common.Message, error) {
    topics, err := r.client.SMembers(r.ctx, retainedIndex).Result()
    if err != nil {
        return nil, err
    }
    
    var matchingTopics []string
    for _, topic := range topics {
        if r.topicMatch(topic, topicFilter) {
            matchingTopics = append(matchingTopics, topic)
        }
    }
    
    if len(matchingTopics) == 0 {
        return nil, nil
    }
    
    // 批量获取匹配的保留消息
    keys := make([]string, len(matchingTopics))
    for i, topic := range matchingTopics {
        keys[i] = retainedPrefix + topic
    }
    
    values, err := r.client.MGet(r.ctx, keys...).Result()
    if err != nil {
        return nil, err
    }
    
    var messages []*common.Message
    for _, value := range values {
        if value == nil {
            continue
        }
        
        var message common.Message
        if err := json.Unmarshal([]byte(value.(string)), &message); err != nil {
            continue
        }
        
        messages = append(messages, &message)
    }
    
    return messages, nil
}

// 主题匹配逻辑（与内存存储保持一致）
func (r *RedisStore) topicMatch(topic, filter string) bool {
    if filter == topic {
        return true
    }
    
    filterParts := strings.Split(filter, "/")
    topicParts := strings.Split(topic, "/")
    
    return r.matchParts(topicParts, filterParts, 0, 0)
}

func (r *RedisStore) matchParts(topic, filter []string, ti, fi int) bool {
    if fi >= len(filter) {
        return ti >= len(topic)
    }
    
    if filter[fi] == "#" {
        return true
    }
    
    if ti >= len(topic) {
        return false
    }
    
    if filter[fi] == "+" || filter[fi] == topic[ti] {
        return r.matchParts(topic, filter, ti+1, fi+1)
    }
    
    return false
}

// 批量操作支持
func (r *RedisStore) SaveMultipleMessages(clientID string, messages []*common.Message) error {
    if len(messages) == 0 {
        return nil
    }
    
    key := messagePrefix + clientID
    pipe := r.client.Pipeline()
    
    for _, message := range messages {
        msgBytes, err := json.Marshal(message)
        if err != nil {
            continue
        }
        
        score := float64(message.Timestamp.UnixNano())
        pipe.ZAdd(r.ctx, key, &redis.Z{Score: score, Member: string(msgBytes)})
    }
    
    pipe.ZRemRangeByRank(r.ctx, key, 0, -1001)
    pipe.Expire(r.ctx, key, 24*time.Hour)
    
    _, err := pipe.Exec(r.ctx)
    return err
}

// 获取统计信息
func (r *RedisStore) GetStats() (map[string]interface{}, error) {
    stats := make(map[string]interface{})
    
    pipe := r.client.Pipeline()
    
    // 统计各种数据
    sessionKeysCmd := pipe.Keys(r.ctx, sessionPrefix+"*")
    retainedCountCmd := pipe.SCard(r.ctx, retainedIndex)
    subscriptionKeysCmd := pipe.Keys(r.ctx, subscriptionPrefix+"*")
    messageKeysCmd := pipe.Keys(r.ctx, messagePrefix+"*")
    
    _, err := pipe.Exec(r.ctx)
    if err != nil {
        return nil, err
    }
    
    sessionKeys, _ := sessionKeysCmd.Result()
    retainedCount, _ := retainedCountCmd.Result()
    subscriptionKeys, _ := subscriptionKeysCmd.Result()
    messageKeys, _ := messageKeysCmd.Result()
    
    stats["sessions"] = len(sessionKeys)
    stats["retained_messages"] = retainedCount
    stats["subscription_sets"] = len(subscriptionKeys)
    stats["message_queues"] = len(messageKeys)
    
    // 计算总消息数
    var totalMessages int64
    if len(messageKeys) > 0 {
        for _, key := range messageKeys {
            count, _ := r.client.ZCard(r.ctx, key).Result()
            totalMessages += count
        }
    }
    stats["total_messages"] = totalMessages
    
    return stats, nil
}

// 清理过期数据
func (r *RedisStore) CleanupExpiredData() error {
    // Redis会自动处理TTL，这里可以做一些额外的清理工作
    // 比如清理孤立的索引数据等
    
    // 清理保留消息索引中的无效条目
    topics, err := r.client.SMembers(r.ctx, retainedIndex).Result()
    if err != nil {
        return err
    }
    
    pipe := r.client.Pipeline()
    for _, topic := range topics {
        key := retainedPrefix + topic
        pipe.Exists(r.ctx, key)
    }
    
    results, err := pipe.Exec(r.ctx)
    if err != nil {
        return err
    }
    
    // 移除不存在的主题
    var toRemove []interface{}
    for i, result := range results {
        if cmd, ok := result.(*redis.IntCmd); ok {
            if cmd.Val() == 0 {
                toRemove = append(toRemove, topics[i])
            }
        }
    }
    
    if len(toRemove) > 0 {
        return r.client.SRem(r.ctx, retainedIndex, toRemove...).Err()
    }
    
    return nil
}

// 健康检查
func (r *RedisStore) Ping() error {
    return r.client.Ping(r.ctx).Err()
}

// 关闭连接
func (r *RedisStore) Close() error {
    return r.client.Close()
}

// 辅助函数
func parsePacketID(s string) uint16 {
    var id uint16
    fmt.Sscanf(s, "%d", &id)
    return id
}
